import torch
from torchvision import transforms
from PIL import Image
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms
from PIL import Image
import pandas as pd
import os
from torchvision import models
# 定义颜值打分器模型
class AttractivenessScorer(nn.Module):
    def __init__(self):
        super(AttractivenessScorer, self).__init__()
        # 使用预训练的EfficientNet作为特征提取器
        self.backbone = models.efficientnet_b0(pretrained=True)
        # 替换最后的分类层
        self.backbone.classifier = nn.Sequential(
            nn.Linear(self.backbone.classifier[1].in_features, 128),
            nn.ReLU(),
            nn.Linear(128, 1),
            nn.Sigmoid()  # 输出在0到1之间
        )
    
    def forward(self, x):
        x = self.backbone(x)
        return x * 10  # 将输出缩放到0到10之间
# 加载模型
model = AttractivenessScorer()
model.load_state_dict(torch.load('attractiveness_scorer.pt'))
model.eval()

# 数据预处理
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# 推理函数
def predict_attractiveness(image_path):
    image = Image.open(image_path).convert('RGB')
    image = transform(image).unsqueeze(0)  # 增加batch维度
    with torch.no_grad():
        output = model(image)
    return output.item()

# 测试推理
image_path = 'test.jpg'
score = predict_attractiveness(image_path)
print(f'Predicted attractiveness score: {score:.2f}')